# -*- coding: utf-8 -*-
#!/usr/local/bin/python
# Author: yuxiang yu

'''
国产向量数据库

安装依赖
pip install pymilvus==2.2.13

milvus 版本
2.2.11

# sort
1. create collection
2. crete index
3. load
4. search
'''
from pymilvus import connections, utility, FieldSchema, CollectionSchema, DataType, Collection, Partition
import datetime
from django.conf import settings


class MilvusDB(object):
    global _ALIAS
    global _HOST
    global _PORT
    global _USER
    global _PASSWORD
    global _INDEX_FILE_SIZE
    global _NLIST
    global _NPROBE
    global index_params_ip
    global index_params_l2

    # 开发环境
    _ALIAS = "default"  # 别名
    _HOST = settings.MILVUS_HOST
    _PORT = settings.MILVUS_PORT
    _USER = settings.MILVUS_USER
    _PASSWORD = settings.MILVUS_PASSWORD

    _INDEX_FILE_SIZE = 1024
    _NLIST = 2048
    _NPROBE = 16
    index_params_ip = {
        "metric_type": "IP",
        "index_type": "IVF_FLAT",
        "params": {"nlist": 128},
    }
    index_params_l2 = {
        "metric_type": "L2",
        "index_type": "IVF_FLAT",
        "params": {"nlist": 128},
    }
    
    def __init__(
        self,
        collection_info
    ):
        self.rq = datetime.datetime.now()
        self.collection_info = collection_info
        self.collection_name_dict = {}
        self.collection_dim_dict = {}
        assert self._conn() is True
        self.load_collection()
        self.set_index()
        self.output_fields = ["code", "question", "source"]

    def load_collection(self):
        try:
            self.load_or_create_collection()
        except:
            # self.load_or_create_collection().release()
            self.load_or_create_collection()

    def __del__(self):
        try:
            # self.collection.release()
            self._disconnect()
        except Exception as e:
            print(f"auto release failed {e}")

    def __repr__(self):
        return f"db [{_HOST}], rq {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"

    def _conn(self):
        if _USER and _PASSWORD:
            connections.connect(
                alias = _ALIAS,
                host = _HOST,
                port = _PORT,
                user=_USER,
                password=_PASSWORD,
                db_name = settings.MILVUS_DB_NAME,  # 默认数据库
            )
        else:
            connections.connect(
                alias=_ALIAS,
                host=_HOST,
                port=_PORT,
                db_name=settings.MILVUS_DB_NAME,  # 默认数据库
            )
        try:
            if connections.has_connection(_ALIAS):
                print("connection success")
                return True
            return False
        except Exception as e:
            print(f"connection error: {e}")
            exit(-1)
            return False

    def _disconnect(self):
        try:
            connections.disconnect(_ALIAS)
        except Exception as e:
            print(f"disconnect error: {e}")

    @property
    def show_all_columns(self):
        return ["code", "question", "feature", "user", "session_id", "embedding_type", "create_time"]

    def set_schema(self, embedding_dim, collection_name):
        fields_list = [
            FieldSchema(name="code", dtype=DataType.VARCHAR, is_primary=True, auto_id=False, max_length=32),
            FieldSchema(name="question", dtype=DataType.VARCHAR, max_length=2048),
            FieldSchema(name="feature", dtype=DataType.FLOAT_VECTOR, dim=embedding_dim),
            FieldSchema(name="user", dtype=DataType.VARCHAR, max_length=64),
            FieldSchema(name="session_id", dtype=DataType.VARCHAR, max_length=64),
            FieldSchema(name="embedding_type", dtype=DataType.VARCHAR, max_length=64),
            FieldSchema(name="source", dtype=DataType.VARCHAR, max_length=64),
            FieldSchema(name="create_time", dtype=DataType.VARCHAR, max_length=50)
        ]
        schema = CollectionSchema(
            fields=fields_list,
            description=f"collection: {collection_name}",
            enable_dynamic_field=True
            )
        return schema

    def load_or_create_collection(self):
        for collection_name, real_collection_name, dim_list in self.collection_info:
            for dim in dim_list:
                if collection_name not in self.collection_dim_dict:
                    self.collection_dim_dict[collection_name] = [dim]
                else:
                    if dim not in self.collection_dim_dict:
                        self.collection_dim_dict[collection_name].append(dim)

                collection_ip_name = collection_name + "_{}_ip".format(dim)
                collection_l2_name = collection_name + "_{}_l2".format(dim)
                real_collection_ip_name = real_collection_name + "_{}_ip".format(dim)
                real_collection_l2_name = real_collection_name + "_{}_l2".format(dim)
                if utility.has_collection(real_collection_ip_name):
                    self.collection_name_dict[collection_ip_name] = Collection(real_collection_ip_name)
                else:
                    print(f"init databse, start create collection {collection_ip_name}")
                    schema_ip = self.set_schema(dim, real_collection_ip_name)
                    collection_ip = Collection(
                            name=real_collection_ip_name,
                            schema=schema_ip,
                            shards_num=2,
                        )
                    self.collection_name_dict[collection_ip_name] = collection_ip
                if utility.has_collection(real_collection_l2_name):
                    self.collection_name_dict[collection_l2_name] = Collection(real_collection_l2_name)
                else:
                    print(f"init databse, start create collection {collection_l2_name}")
                    schema_l2 = self.set_schema(dim, real_collection_l2_name)
                    collection_l2 = Collection(
                            name=real_collection_l2_name,
                            schema=schema_l2,
                            shards_num=2,
                        )
                    self.collection_name_dict[collection_l2_name] = collection_l2

    
    def build_collection(self, collection_name, real_collection_name, embedding_dim):
        collection_ip_name = collection_name + "_{}_ip".format(embedding_dim)
        collection_l2_name = collection_name + "_{}_l2".format(embedding_dim)
        real_collection_ip_name = real_collection_name + "_{}_ip".format(embedding_dim)
        real_collection_l2_name = real_collection_name + "_{}_l2".format(embedding_dim)
        if collection_name not in self.collection_dim_dict:
            self.collection_dim_dict[collection_name] = [embedding_dim]
        else:
            if embedding_dim not in self.collection_dim_dict:
                self.collection_dim_dict[collection_name].append(embedding_dim)

        if not utility.has_collection(real_collection_ip_name):
            schema_ip = self.set_schema(embedding_dim, real_collection_ip_name)
            collection_ip = Collection(
                    name=real_collection_ip_name,
                    schema=schema_ip,
                    shards_num=2,
                )
            self.collection_name_dict[collection_ip_name] = collection_ip

            schema_l2 = self.set_schema(embedding_dim, real_collection_l2_name)
            collection_l2 = Collection(
                    name=real_collection_l2_name,
                    schema=schema_l2,
                    shards_num=2,
                )
            self.collection_name_dict[collection_l2_name] = collection_l2
            print(f"init databse, start create collection {collection_ip_name} and {collection_l2_name}")
            vector_field_name = "feature"
            self.collection_name_dict[collection_ip_name].create_index(field_name=vector_field_name, index_params=index_params_ip)
            self.collection_name_dict[collection_ip_name].load()
            self.collection_name_dict[collection_l2_name].create_index(field_name=vector_field_name, index_params=index_params_l2)
            self.collection_name_dict[collection_l2_name].load()
            print(f"set index and load finish!")
            msg = f"collection {collection_name} create success!"
        else:
            msg = f"collection {collection_name} exists!"
        return True, msg
    
    def del_collection(self, collection_name, real_collection_name):
        dim_list = self.collection_dim_dict[collection_name]
        for embedding_dim in dim_list:
            collection_ip_name = collection_name + "_{}_ip".format(embedding_dim)
            collection_l2_name = collection_name + "_{}_l2".format(embedding_dim)
            real_collection_ip_name = real_collection_name + "_{}_ip".format(embedding_dim)
            real_collection_l2_name = real_collection_name + "_{}_l2".format(embedding_dim)

            utility.drop_collection(real_collection_ip_name)
            self.collection_name_dict.pop(collection_ip_name)

            utility.drop_collection(real_collection_l2_name)
            self.collection_name_dict.pop(collection_l2_name)
        self.collection_dim_dict.pop(collection_name)
        return True

    
    def del_partition(self, collection_name, partition_name, embedding_dim):
        # dim_list = self.collection_dim_dict[collection_name]
        # for embedding_dim in dim_list:
        collection_name_ip = collection_name + "_{}_ip".format(embedding_dim)
        collection_name_l2 = collection_name + "_{}_l2".format(embedding_dim)
        try:
            if self.collection_name_dict[collection_name_ip].has_partition(partition_name=partition_name):
                self.collection_name_dict[collection_name_ip].release()
                partition_ip = Partition(self.collection_name_dict[collection_name_ip], partition_name)
                partition_ip.release()
                self.collection_name_dict[collection_name_ip].drop_partition(partition_name)
                self.collection_name_dict[collection_name_ip].load()

            if self.collection_name_dict[collection_name_l2].has_partition(partition_name=partition_name):
                self.collection_name_dict[collection_name_l2].release()
                partition_l2 = Partition(self.collection_name_dict[collection_name_l2], partition_name)
                partition_l2.release()
                self.collection_name_dict[collection_name_l2].drop_partition(partition_name)
                self.collection_name_dict[collection_name_l2].load()
        except Exception as e:
            print(f"drop partition error: {e}")
            return False
        return True
    
    def rename_collection(self, old_collection_name, new_collection_name):
        dim_list = self.collection_dim_dict[old_collection_name]
        for embedding_dim in dim_list:
            old_collection_name_ip = old_collection_name + "_{}_ip".format(embedding_dim)
            old_collection_name_l2 = old_collection_name + "_{}_l2".format(embedding_dim)
            
            new_collection_name_ip = new_collection_name + "_{}_ip".format(embedding_dim)
            new_collection_name_l2 = new_collection_name + "_{}_l2".format(embedding_dim)        

            self.collection_name_dict[new_collection_name_ip] = self.collection_name_dict.pop(old_collection_name_ip)
            self.collection_name_dict[new_collection_name_l2] = self.collection_name_dict.pop(old_collection_name_l2)
        self.collection_dim_dict[new_collection_name] = self.collection_dim_dict.pop(old_collection_name)
        return True
    
    def check_collection(self, collection):
        # 检查 collection
        desc = f"""schema: {collection.schema}\ndescription: {collection.description}\nname: {collection.name}\nis_empty: {collection.is_empty}\nnum_entities: {collection.num_entities}\nprimary_field: {collection.primary_field}\npartitions: {collection.partitions}\nindexes: {collection.indexes}\n"""
        print(desc)

    def get_num_entities(self, collection):
        return collection.num_entities

    def get_primary_field(self, collection):
        return collection.primary_field

    def get_partitions(self, collection):
        return collection.partitions

    def list_collection(self):
        try:
            # self.collection.load()
            return utility.list_collections()
        except Exception as e:
            print(f"list_collection: {e}")
            return None

    def list_partition(self, collection):
        try:
            return collection.partitions
        except Exception as e:
            print(f"list partition error: {e}")
            return False

    def drop_partition(self, partition_name):
        try:
            self.collection_ip.drop_partition(partition_name)
            self.collection_l2.drop_partition(partition_name)
            return True
        except Exception as e:
            print(f"drop partition error: {e}")
            self.collection_ip.release()
            self.collection_ip.drop_partition(partition_name)
            self.collection_l2.release()
            self.collection_l2.drop_partition(partition_name)
            return False

#     def load_partition(self, partition_name):
#         if isinstance(partition_name, str):
#             partition_name = [str(partition_name).strip()]

#         if not isinstance(partition_name, list):
#             partition_name = [partition_name]

#         try:
#             self.collection_ip.load(partition_name, replica_number=2)
#             self.collection_l2.load(partition_name, replica_number=2)
#             return True
#         except Exception as e:
#             print(f"load partition error: {e}")
#             return False

#     def release_collection(self):
#         try:
#             self.collection_ip.release()
#             self.collection_l2.release()
#             return True
#         except Exception as e:
#             print(f"release collection error: {e}")
#             return False

#     def release_partition(self, partition_name):
#         if isinstance(partition_name, str):
#             partition_name = str(partition_name).strip()
#             partition = Partition(self.collection, partition_name)
#             partition.release()
#             return True

#         if isinstance(partition_name, (tuple, list)):
#             for name in partition_name:
#                 try:
#                     partition = Partition(self.collection, name)
#                     partition.release()
#                 except Exception as e:
#                     print(f"release partition error: {e}")
#             return True
#         print(f"partition_name: {partition_name}")
#         return False

    # def delete_entity(self, expr):
    #     # expr = "sku in [0,1]"
    #     try:
    #         self.collection_ip.delete(expr)
    #         self.collection_l2.delete(expr)
    #         return True
    #     except Exception as e:
    #         print(f"delete entity error: {e}")
    #         return False

    def set_partition(self, collection, partition_name=None, timeout=20):
        # TODO: create a partition
        if not collection.has_partition(
                partition_name=partition_name,
                timeout=timeout):
            try:
                if partition_name is None or str(partition_name).startswith("None"):
                    partition_name = "part1"
                collection.create_partition(
                    partition_name=partition_name,
                    description=partition_name
                    )
                return True
            except Exception as e:
                print(f"set partition error:{e}, rq: {self.rq}")
        print(f"set_partition: {partition_name} 已存在, rq:{self.rq}")
        return True

    def set_index(self):
        for collection_name, collection in self.collection_name_dict.items():
            vector_field_name = "feature"
            if not collection.has_index():
                if collection_name.endswith("ip"):
                    collection.create_index(field_name=vector_field_name, index_params=index_params_ip)
                if collection_name.endswith("l2"):
                    collection.create_index(field_name=vector_field_name, index_params=index_params_l2)
                collection.load()

        # for collection_name, real_collection_name, dim_list in self.collection_info:
        #     dim_list = self.collection_dim_dict[collection_name]
        #     for embedding_dim in dim_list:
        #         collection_ip_name = collection_name + "_{}_ip".format(embedding_dim)
        #         collection_l2_name = collection_name + "_{}_l2".format(embedding_dim)
        #         vector_field_name = "feature"
        #         if not self.collection_name_dict[collection_ip_name].has_index():
        #             self.collection_name_dict[collection_ip_name].create_index(field_name=vector_field_name, index_params=index_params_ip)
        #             self.collection_name_dict[collection_ip_name].load()
        #         if not self.collection_name_dict[collection_l2_name].has_index():
        #             self.collection_name_dict[collection_l2_name].create_index(field_name=vector_field_name, index_params=index_params_l2)
        #             self.collection_name_dict[collection_l2_name].load()

    def delete_data(self, del_idx_dict, collection_name: str):
        try:
            # 首列唯一索引
            for embedding_dim, del_idx_list in del_idx_dict.items():
                expr = f"code in {del_idx_list}"
                self.collection_name_dict[collection_name + "_{}_ip".format(embedding_dim)].delete(expr)
                self.collection_name_dict[collection_name + "_{}_l2".format(embedding_dim)].delete(expr)
            return True
        except Exception as e:
            print(f"delete data error {e}")
            return False
    
    def insert_data(self, data, collection_name: str, partition_name: str, embedding_dim: str, timeout=30):
        if not connections.has_connection(_ALIAS):
            self._conn()
            self.load_collection()
            self.set_index()
        self.set_partition(self.collection_name_dict[collection_name+"_{}_ip".format(embedding_dim)], partition_name=partition_name, timeout=timeout)
        self.set_partition(self.collection_name_dict[collection_name+"_{}_l2".format(embedding_dim)], partition_name=partition_name, timeout=timeout)

        init_ip_insert_status = self.collection_name_dict[collection_name+"_{}_ip".format(embedding_dim)].insert(data, partition_name=partition_name)
        init_l2_insert_status = self.collection_name_dict[collection_name+"_{}_l2".format(embedding_dim)].insert(data, partition_name=partition_name)
        return True, "success"


    def _result_dict(self, hit):
        tmp_dict = {
            "code": hit.entity.get("code"),
            "question": hit.entity.get("question"),
            "source": hit.entity.get("source"),
            "score": round(hit.distance, 3)
            }
        return tmp_dict

    def get_hybrid_vector_search(self, query_vector, session_id, embedding_type, embedding_dim, collection_name, partition_name, similar_type, topk=1, timeout=10, nprobe=10, is_dict=True):
        if similar_type == "IP":
            similar_param = {
                "metric_type": "IP",
                "params": {"nprobe": nprobe},
            }
        elif similar_type == "L2":
            similar_param = {
                "metric_type": "L2",
                "params": {"nprobe": nprobe},
            }
        else:
            similar_param = {
                "metric_type": "IP",
                "params": {"nprobe": nprobe},
            }
        if session_id is None:
            expr = f"embedding_type == '{embedding_type}'"
        else:
            expr = f"embedding_type == '{embedding_type}' and session_id == '{session_id}'"
        partition_names = [partition_name]
        # 范围检索
        search_param = {
            "data": query_vector,
            "anns_field": "feature",
            "param": similar_param,
            "offset": 0,
            "limit": int(topk),
            "partition_names": partition_names,
            "expr": expr,
            "output_fields": self.output_fields,
            "timeout": timeout,
            "consistency_level": "Strong"
        }
        if not connections.has_connection(_ALIAS):
            self._conn()
            self.load_collection()
            self.set_index()
        self.set_partition(self.collection_name_dict[
            collection_name+"_{}_ip".format(embedding_dim)], partition_name=partition_name, timeout=timeout)
        self.set_partition(self.collection_name_dict[
            collection_name+"_{}_l2".format(embedding_dim)], partition_name=partition_name, timeout=timeout)
        try:
            if similar_type == "IP":
                result = self.collection_name_dict[collection_name+"_{}_ip".format(embedding_dim)].search(**search_param)
            elif similar_type == "L2":
                result = self.collection_name_dict[collection_name+"_{}_l2".format(embedding_dim)].search(**search_param)
            else:
                result = self.collection_name_dict[collection_name+"_{}_ip".format(embedding_dim)].search(**search_param)
        except:
            return None
            # self.load_collection()
            # if similar_type == "IP":
            #     result = self.collection_name_dict[collection_name+"_{}_ip".format(embedding_dim)].search(**search_param)
            # elif similar_type == "L2":
            #     result = self.collection_name_dict[collection_name+"_{}_l2".format(embedding_dim)].search(**search_param)
            # else:
            #     result = self.collection_name_dict[collection_name+"_{}_ip".format(embedding_dim)].search(**search_param)
        if is_dict:
            res_list = []
            for hits in result:
                res_list.append(list(map(self._result_dict, list(hits))))
            return res_list[0]

        return result


if __name__ == "__main__":
    import numpy as np
    # 初始化
    vector_db = MilvusDB(collection_name_list=["dialogue_test1", "dialogue_test2"])
    # 写入数据库
    # tmp_result_list = []
    # code_list = [f"{i}_1" for i in range(5)]
    # question_list = [f"question_{i}" for i in range(5)]
    # feature_list = [np.random.random((1, 768)).tolist()[0] for _ in range(5)]
    # user_list = ["" for i in range(5)]
    # session_id_list = ["124" for i in range(5)]
    # embedding_type_list = ["m3e" for i in range(5)]
    # source_list = ["test" for i in range(5)]
    # status_list = [1 for _ in range(5)]
    # create_time_list = [datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') for _ in range(5)]
    # tmp_result_list.extend([code_list, question_list, feature_list, user_list, session_id_list, embedding_type_list, project_name_list, source_list, status_list, create_time_list])
    # vector_db.insert_data(collection_name="dialogue_test2", data=tmp_result_list, partition_name='test')
    # 创建数据库
    field_type_dict = {
        "str": DataType.VARCHAR,
        "vector": DataType.FLOAT_VECTOR,
        "int": DataType.INT64
    }
    return_status, msg = vector_db.build_collection("test", [{"name": "code", "type": "str", "length": 64},
                                       {"name": "question", "type": "str", "length": 1024},
                                       {"name": "embedding", "type": "vector", "length": 768},
                                       {"name": "status", "type": "int", "length": 2}], field_type_dict)
    print(return_status, msg)
    # print(vector_db.get_hybrid_vector_search(self, query_vector, embedding_type, project_name, session_id, similar_type, topk=1, timeout=10, nprobe=10, is_dict=True))
    # vector_db.release_collection()
    # print(status)
#     from uuid import uuid4
#     import numpy as np
#     import datetime
#     nums = 3

#     # TODO: 写入数据
#     db = MilvusDB()
#     # 显示所有列
#     print(db.show_all_columns)
#     # db.del_collection()
#     # 模拟数据
#     tmp_result_list = []
    
#     ids = [i for i in range(3, 6)]
#     question = [f"question_{i}" for i in range(3, 6)]
#     feature = [np.random.random((1, 768)).tolist()[0] for _ in range(nums)]
#     status = [str(1) for _ in range(nums)]
#     create_time = [datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S') for _ in range(nums)]
#     tmp_result_list.extend([ids, question, feature, status, create_time])
#     # 写入数据
#     db.insert_data(data=tmp_result_list, partition_name='test')
    # 手动关闭
    # db.manual_close()
    # del db

    # TODO: 检索数据
    # 检索
#     db_search = MilvusDB(only_search=True)
#     result = db_search.get_search([feature[0]], partition_names=['test', '_default'], topk=10, nprobe=10, is_dict=True)
#     print(f"识别结果", result)

#     result = db_search.get_search([feature[1]], partition_names=['test', '_default'], topk=10, nprobe=10, is_dict=True)
#     print(f"识别结果", result)
#     # 手动释放
#     db_search.manual_close()
#     del db_search

    # TODO: 其他指标
    # print("=== 其他指标 ===")
    # db_info = MilvusDB()
    # print(db_info.get_num_entities())
    # db_info.manual_close()
    # del db_info

